package main

type TreeNode struct {
	Val   int
	Left  *TreeNode
	Right *TreeNode
}

func widthOfBinaryTree(root *TreeNode) int {
	que := []*TreeNode{root}
	indexMap := map[*TreeNode]int{
		root: 0,
	}
	depthMap := map[*TreeNode]int{
		root: 0,
	}
	minIndexMap := make(map[int]int)
	maxIndexMap := make(map[int]int)
	for len(que) > 0 {
		node := que[0]
		que = que[1:]
		idx := indexMap[node]
		depth := depthMap[node]
		if node.Left != nil {
			que = append(que, node.Left)
			indexMap[node.Left] = 2 * idx
			depthMap[node.Left] = depth + 1
			if _, ok := minIndexMap[depth+1]; !ok || minIndexMap[depth+1] > 2*idx {
				minIndexMap[depth+1] = 2 * idx
			}
			if _, ok := maxIndexMap[depth+1]; !ok || maxIndexMap[depth+1] < 2*idx {
				maxIndexMap[depth+1] = 2 * idx
			}
		}
		if node.Right != nil {
			que = append(que, node.Right)
			indexMap[node.Right] = 2*idx + 1
			depthMap[node.Right] = depth + 1
			if _, ok := minIndexMap[depth+1]; !ok || minIndexMap[depth+1] > 2*idx+1 {
				minIndexMap[depth+1] = 2*idx + 1
			}
			if _, ok := maxIndexMap[depth+1]; !ok || maxIndexMap[depth+1] < 2*idx+1 {
				maxIndexMap[depth+1] = 2*idx + 1
			}
		}
	}
	maxWidth := 0
	for depth, min := range minIndexMap {
		max := maxIndexMap[depth]
		width := max - min
		if width > maxWidth {
			maxWidth = width
		}
	}
	return maxWidth + 1
}

func main() {
	n0 := &TreeNode{Val: 1}
	n1 := &TreeNode{Val: 2}
	n2 := &TreeNode{Val: 3}
	n3 := &TreeNode{Val: 4}
	//n4 := &TreeNode{Val: 5}
	//n5 := &TreeNode{Val: 6}
	//n6 := &TreeNode{Val: 7}
	n0.Left = n1
	n0.Right = n2
	n1.Left = n3
	//n2.Right = n4
	//n3.Left = n5
	//n4.Left = n6
	println(widthOfBinaryTree(n0))
}
